# python image library python图像处理库
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
import torch.nn as nn
import torch

from Lesson.MLP.lession10_MyNet import MyNet

net = MyNet(28*28*1, 300, 10)
# #保存模型
# torch.save(net, 'hpu_net.pth')
# #加载模型
# torch.load('hpu_net.pth')
# #得到模型中的参数
# list(net.named_parameters())
#
# # 模型的所有参数，仅有参数
# net.state_dict()
#
# net_3 = MyNet()
#
# torch.save(net.state_dict(), 'my_net.pth')
#
# #加载
# net_3.load_state_dict(torch.load('my_net.pth'))

p = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
test = datasets.MNIST('../data', transform=p, train=False, download=True)
test_loader = DataLoader(test, batch_size=16, shuffle=False)
loss_fun = nn.CrossEntropyLoss()
# 测试模型
eval_loss = 0.0
eval_acces = 0.0
with torch.no_grad():
    for img,label in test_loader:
        out = net(img.view(img.shape[0], -1))
        loss = loss_fun(out, label)
        eval_loss += loss.item()
        # pred = out.argmax(dim=1)
        _, pred = out.max(dim=1)
        num_correct = (pred == label).sum().item()
        acc = num_correct/img.shape[0]
        eval_acces += acc
    print(f'test_loss:{eval_loss/len(test_loader):.4f}, test_acc:{eval_acces/len(test_loader)}')



